import matplotlib.pyplot as plt
import numpy as np

keypoint_names = [
    'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
    'left_wrist', 'right_wrist', 'left_hip', 'right_hip',
    'left_knee', 'right_knee', 'left_ankle', 'right_ankle'
]

skeleton = [
    [15, 13], [13, 11], [16, 14], [14, 12],
    [11, 12], [5, 11], [6, 12], [5, 6],
    [5, 7], [6, 8], [7, 9], [8, 10],
    [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],
    [0, 5], [0, 6]
]

keypoints = [
    (300, 100),    # nose
    (290, 90), (310, 90),        # eyes
    (280, 100), (320, 100),      # ears
    (260, 180), (340, 180),      # shoulders
    (240, 260), (360, 230),      # elbows
    (220, 340), (380, 270),      # wrists
    (270, 360), (330, 360),      # hips
    (250, 470), (345, 420),      # knees
    (240, 580), (360, 500),      # ankles
]

left_indices = [1, 3, 5, 7, 9, 11, 13, 15]
right_indices = [2, 4, 6, 8, 10, 12, 14, 16]

def draw_pose_colored(keypoints, skeleton, keypoint_names):
    plt.figure(figsize=(6, 10))
    ax = plt.gca()

    cmap = plt.cm.rainbow
    colors = [cmap(i / len(skeleton)) for i in range(len(skeleton))]

    for i, joint_pair in enumerate(skeleton):
        pt1 = keypoints[joint_pair[0]]
        pt2 = keypoints[joint_pair[1]]
        color = colors[i]
        plt.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]], '-', linewidth=4, color=color)

    for i, (x, y) in enumerate(keypoints):
        plt.plot(x, y, 'o', color='black', markersize=4)
        if i in left_indices:
            plt.text(x - 10, y, keypoint_names[i], fontsize=8, color='black', ha='right')
        elif i in right_indices:
            plt.text(x + 10, y, keypoint_names[i], fontsize=8, color='black', ha='left')
        else:
            plt.text(x, y + 10, keypoint_names[i], fontsize=8, color='black', ha='left')

    ax.set_xlim(150, 450)
    ax.set_ylim(620, 50)
    ax.set_aspect('equal')
    plt.axis('off')
    plt.show()


draw_pose_colored(keypoints, skeleton, keypoint_names)
